import numpy as np
from PIL import Image
import torch
import cv2
from torchvision.transforms import functional as F


class ImageToTensor(object):
    '''
    将PIL图像或者nparray转换为tensor
    torchvision中的ToTensor无法将单通道的nparray转换为tensor，这里添加这个功能。
    '''
    def __init__(self):
        pass
    def __call__(self, pic):
        # 如果pic是单通道的nparray，直接使用from_numpy函数转换为tensor
        if isinstance(pic, np.ndarray) and len(pic.shape) == 2:
            img = torch.from_numpy(pic)
            # backward compatibility
            if isinstance(img, torch.ByteTensor):
                return img.float().div(255)
            else:
                return img
        else:
            return F.to_tensor(pic)

class ResizeImage(object):
    """
    Input an numpy array or a PIL image and return a PIL image with given size "new_size", keeping num of images unchanged.
    """
    def __init__(self, new_size):
        self.new_size = new_size
    def __call__(self, image):
        image = cv2.resize(np.asarray(image), self.new_size)
        return Image.fromarray(image)


class VectorLabelNoise(object):
    '''
    给标签向量添加一个小小的扰动，防止损失函数失去最小值点
    '''
    def __init__(self, eplon=0.01):
        self.eplon = eplon
    def __call__(self, label_oh):
        return (1 - self.eplon) * label_oh + self.eplon / len(label_oh)
